#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <inttypes.h>

int main(int argc, char *argv[])
{
    FILE *fp;
    uint8_t *hdcopy;
    uint8_t *plain;
    uint32_t length;
    uint32_t fit_size = 0;

    int i, j, k, r;

    if(argc != 3 && argc != 4)
    {
        fprintf(stderr, "Usage: %s <HD-COPY Image> <IMA plain floppy image> [size]\n\n", argv[0]);
        fprintf(stderr, "This small tool can decompress HD-COPY image to plain floppy image.\n");
        fprintf(stderr, "It can help you to use old HD-COPY image on modern PC,\nespecially on Virtual Machine to test out old-school softwares.\n\n");
        fprintf(stderr, "Zhiyuan Wan <h@iloli.bid> 2018, License WTFPL.\nAlgorithm analyze from <https://github.com/ciel-yu/devnotes>. Thanks him!\n");
        return -1;
    }
    if(argc == 4)
    {
        char *endptr = NULL;
        unsigned long parsed = strtoul(argv[3], &endptr, 10);
        if(endptr == argv[3] || *endptr != '\0' || parsed == 0)
        {
            fprintf(stderr, "Invalid output size parameter: %s\n", argv[3]);
            return -1;
        }
        if(parsed > 0x168000UL)
        {
            fprintf(stderr, "Output size too large (max 0x168000): %lu\n", parsed);
            return -1;
        }
        fit_size = (uint32_t)parsed;
    }
    fp = fopen(argv[1], "rb");
    if(!fp)
    {
        fprintf(stderr, "Can't open source HD-COPY image!\n");
        return -1;
    }
    if(fseek(fp, 0, SEEK_END) != 0)
    {
        fprintf(stderr, "Failed to seek to end of input file.\n");
        fclose(fp);
        return -1;
    }
    long flen = ftell(fp);
    if(flen <= 0)
    {
        fprintf(stderr, "Invalid input file size.\n");
        fclose(fp);
        return -1;
    }
    length = (uint32_t)flen;

    hdcopy = (uint8_t*)malloc(length);
    plain = (uint8_t*)malloc(0x168000);
    if(!hdcopy || !plain)
    {
        fprintf(stderr, "Memory allocation failed.\n");
        fclose(fp);
        free(hdcopy);
        free(plain);
        return -1;
    }

    if(fseek(fp, 0, SEEK_SET) != 0)
    {
        fprintf(stderr, "Failed to seek to start of input file.\n");
        fclose(fp);
        free(hdcopy);
        free(plain);
        return -1;
    }

    size_t rd = fread(hdcopy, 1, length, fp);
    if(rd != length)
    {
        fprintf(stderr, "Reading input file failed (read %zu of %u bytes).\n", rd, length);
        fclose(fp);
        free(hdcopy);
        free(plain);
        return -1;
    }

    printf("Decompressing HD-COPY Image\nInput size = %u\n", length);
    fclose(fp);

    uint8_t *actualimage;
    uint8_t *payload;

    int is20 = (length >= 2 && hdcopy[0] == 0xff && hdcopy[1] == 0x18);

    if(is20)
    {
        printf("That is an HD-COPY 2.0 Image\n");
        if(length < 0x0e + 2 + 168)
        {
            fprintf(stderr, "Corrupt or truncated HD-COPY 2.0 header.\n");
            goto deal;
        }
        actualimage = hdcopy + 0x0e; /* 跳过标有卷标的文件头 */
        payload = actualimage + 2 + 168; /* 载荷段 */
    }
    else
    {
        printf("That is an HD-COPY 1.7 Image\n");
        if(length < 2 + 164)
        {
            fprintf(stderr, "Corrupt or truncated HD-COPY 1.7 header.\n");
            goto deal;
        }
        actualimage = hdcopy;
        payload = actualimage + 2 + 164;
    }
    /* 开始解码 */
    int maxTrackCount = actualimage[0];
    int secPerTrack = actualimage[1];
    int heads = 2;

    if(secPerTrack <= 0 || secPerTrack > 36)
    {
        fprintf(stderr, "Invalid sectors per track: %d\n", secPerTrack);
        goto deal;
    }
    if(maxTrackCount < 0 || maxTrackCount > 84)
    {
        fprintf(stderr, "Invalid max track count: %d\n", maxTrackCount);
        goto deal;
    }

    printf("maxTrackCount = %d, secPerTrack = %d\n", maxTrackCount + 1, secPerTrack);
    uint8_t *dataTrack = actualimage + 2; /* 有数据的磁道表 */

    uint8_t *decomp_p = plain;
    uint8_t *plain_end = plain + 0x168000;

    memset(plain, 0x00, 0x168000);

    size_t bytes_per_track = (size_t)secPerTrack * 512;
    size_t expected_total_bytes = (size_t)(maxTrackCount + 1) * heads * bytes_per_track;
    if(expected_total_bytes > (size_t)0x168000)
    {
        fprintf(stderr, "Expected output too large: %zu bytes\n", expected_total_bytes);
        goto deal;
    }

    if(payload > hdcopy + length)
    {
        fprintf(stderr, "Invalid payload pointer beyond input.\n");
        goto deal;
    }

    for(i = 0; i <= maxTrackCount; i++)
    {
        for(j = 0; j < heads; j++)
        {
            if(dataTrack[(i * heads) + j] != 0x01)
            {
                if(decomp_p + bytes_per_track > plain_end)
                {
                    fprintf(stderr, "Output overflow while skipping track %d head %d.\n", i, j);
                    goto deal;
                }
                decomp_p += bytes_per_track;
                continue;
            }
            if((size_t)(hdcopy + length - payload) < 2)
            {
                fprintf(stderr, "Truncated payload when reading block length at track %d head %d.\n", i, j);
                goto deal;
            }
            int dataLen = payload[0] + (payload[1] << 8);
            payload += 2;
            if((size_t)(hdcopy + length - payload) < (size_t)dataLen)
            {
                fprintf(stderr, "Truncated payload block at track %d head %d (need %d bytes).\n", i, j, dataLen);
                goto deal;
            }
            if(dataLen < 1)
            {
                /* 空块，无数据 */
                payload += dataLen;
                continue;
            }

            uint8_t escByte; /* RLE 压缩 */
            escByte = payload[0];
            for(k = 1; k < dataLen; k++)
            {
                if(payload[k] == escByte)
                {
                    if(k + 2 >= dataLen)
                    {
                        fprintf(stderr, "Truncated RLE sequence at track %d head %d.\n", i, j);
                        goto deal;
                    }
                    k++;
                    uint8_t repeatByte = payload[k++];
                    int repeat = payload[k];

                    if(decomp_p + repeat > plain_end)
                    {
                        fprintf(stderr, "Output overflow on repeat (%d) at track %d head %d.\n", repeat, i, j);
                        goto deal;
                    }
                    memset(decomp_p, repeatByte, (size_t)repeat);
                    decomp_p += repeat;
                }
                else
                {
                    if(decomp_p + 1 > plain_end)
                    {
                        fprintf(stderr, "Output overflow on write at track %d head %d.\n", i, j);
                        goto deal;
                    }
                    *(decomp_p++) = payload[k];
                }
            }
            payload += dataLen;
        }
    }
    uint16_t secCount16 = (uint16_t)plain[0x13] + ((uint16_t)plain[0x14] << 8);
    uint32_t secCount32 = (uint32_t)plain[0x20] + ((uint32_t)plain[0x21] << 8) + ((uint32_t)plain[0x22] << 16) + ((uint32_t)plain[0x23] << 24);
    uint32_t secCount = secCount16 ? secCount16 : secCount32;
    if(secCount == 0)
    {
        secCount = (uint32_t)((maxTrackCount + 1) * heads * secPerTrack);
    }

    uint32_t write_bytes = fit_size > 0 ? fit_size : (secCount * 512);
    if(write_bytes == 0 || write_bytes > 0x168000)
    {
        write_bytes = (uint32_t)expected_total_bytes;
    }

    printf("Floppy sector count = %u, fitting to %u bytes\n", secCount, write_bytes);
    printf("Decompress operation completed, write it to file\n");
    fp = fopen(argv[2], "wb");
    if(!fp)
    {
        fprintf(stderr, "Can't save plain floppy image!\n");
        goto deal;
    }
    if(fseek(fp, 0, SEEK_SET) != 0)
    {
        fprintf(stderr, "Failed to seek in output file.\n");
        fclose(fp);
        goto deal;
    }
    size_t wr = fwrite(plain, 1, write_bytes, fp);
    if(wr != write_bytes)
    {
        fprintf(stderr, "Write failed (wrote %zu of %u bytes).\n", wr, write_bytes);
        fclose(fp);
        goto deal;
    }
    fclose(fp);

deal:
    free(hdcopy);
    free(plain);
    return 0;
}